#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@File    ：vector.py
@Author  ：平
@Date    ：2025/9/28 14:07 
"""
import time

import aiohttp
from fastapi import APIRouter, Request

from app.config.config import settings
from app.graph.graph import graph
from app.repositry.book_repo import edit_book
from app.schema.base_response import BaseResponse
from app.schema.vecotr import VectorStoreRequest, VectorStoreMetadata
from app.util import file_util
from app.util.file_util import get_relative_file_path
from app.vector.pipline import Pipline
import logging
from app.vector.vector import vectorstore

router = APIRouter(prefix="/vector")
logger = logging.getLogger(__name__)


@router.post("/store", response_model=BaseResponse)
async def store(request: Request, request_body: VectorStoreRequest):
    start_time = time.time() * 1000
    file_path = request_body.bookFile.objName
    try:
        logger.info("进行智能补全")
        # AI智能补全
        state = {
            "is_fill": True,
            "book": request_body.model_dump()
        }
        response = await graph.ainvoke(state)
        book = edit_book(request_body.id, response.get("book", {}))
        request_body.description = book.description
        request_body.category = book.category
        request_body.tags = book.tags
        request_body.summary = book.summary
        logger.info(f"智能补全完成耗时:{time.time() * 1000 - start_time:.2f}ms，进行向量存储")
        # 下载文件
        p = get_relative_file_path(file_path)
        async with aiohttp.ClientSession() as session:
            payload = {
                "objName": request_body.bookFile.objName
            }
            headers = {
                "Authorization": request.headers.get("Authorization")
            }
            async with session.post(settings.BACKEND_SERVICE_URL + "/attachment/download", json=payload,
                                    headers=headers) as response:
                p.write_bytes(await response.read())

        if p.suffix not in Pipline.FileSuffix:
            raise Exception(f"不支持的文件格式{p.suffix}")
        file_suffix = Pipline.FileSuffix(p.suffix)
        metadata = VectorStoreMetadata(**request_body.model_dump())
        docs = Pipline(p, file_suffix, metadata.model_dump()).load().split().get()[1]
        success = vectorstore.store(docs)
        if not success:
            raise Exception("存储失败")
        return BaseResponse(message="成功")
    except Exception:
        raise
    finally:
        file_util.delete(file_path)
        logger.info(f"AI智能检索完成总耗时:{time.time() * 1000 - start_time:.2f}ms")

@router.delete("/remove/{book_id}", response_model=BaseResponse)
async def remove(book_id: int):
    try:
        success = vectorstore.remove(book_id)
        if not success:
            raise Exception("删除失败")
        return BaseResponse(message="成功")
    except Exception:
        raise
